# planner_service.py
import math
import sys
from typing import List, Dict, Tuple
from itertools import groupby
from datetime import datetime

from algorithm.area_coverage_planner import AreaCoveragePlanner
from algorithm.coord_transform import PositionConvert
from algorithm.Astar import AStarPlanner
from algorithm.spatial_Astar import multi_uav_planning, ReservationTable, SLACK
    # 导入Pydantic模型，因为算法内部可能需要引用
from FastAPI.models import PlannerConfig
from FastAPI.logger_config import get_logger

# 获取日志器
logger = get_logger()

class MultiUAVPlanner:
    """多无人机任务规划与路径生成库"""

    def __init__(self, config: PlannerConfig):
        """
        初始化规划器。
        所有可配置参数都从config对象中获取。
        """
        self.config = config  # 保存完整的配置对象
        self.pos_converter = PositionConvert()
        self.ref_lat = self.config.ref_lat
        self.ref_lon = self.config.ref_lon

        # 内部算法常量/默认值
        self.patrol_speed = 6.0
        self.min_patrol_speed = 2.0
        self.max_patrol_speed = 25.0
        self.time_buffer = 120.0
        self.default_altitude = 10.0
        self.sweep_width = 12.0
        self.max_segment_length = 20.0
        self.speed = 8.0  # 初始规划速度
        self.min_planning_speed = 2.0
        self.max_planning_speed = 25.0
        self.speed_increase_factor = 1.1
        self.max_t = 3000

        # 初始化路径规划器实例
        self._init_planners_from_config()
        
        # 状态存储，每次规划前会重置
        self._reset_internal_state()

    def _init_planners_from_config(self):
        """使用配置对象初始化A*路径规划器"""
        min_x, min_y, max_x, max_y = self.config.map_bounds
        obs_x = [min_x, max_x, min_x, max_x]
        obs_y = [min_y, min_y, max_y, max_y]
        
        self.coarse_planner = AStarPlanner(obs_x, obs_y, self.config.coarse_grid_size, self.config.robot_radius)
        self.fine_planner = AStarPlanner(obs_x, obs_y, self.config.fine_grid_size, self.config.robot_radius)
    
    def _reset_internal_state(self):
        """重置所有内部状态以便处理新的任务集"""
        self.uav_positions = {}
        self.uav_goals = {}
        self.task_time_constraints = {}
        self.cruise_data = {}
        self.all_time_paths = {}
        self.patrol_grid_paths = {}
        self.patrol_entry_points = {}
        self.animation_timestamps = []
        self.final_planning_speeds = {}
        self.final_patrol_speeds = {}
        self.reservation_table = ReservationTable()

    def process_tasks(self, tasks_json: List[Dict]) -> Dict:
        """
        处理任务列表，按优先级分组进行统一规划，以确保组内协同避障。
        """
        self._reset_internal_state()

        sorted_tasks = sorted(tasks_json, key=lambda x: x['priority'])

        for task in sorted_tasks:
            self._process_start_positions(task['start'])
            for uav_id in task['uav_ids']:
                self.task_time_constraints[uav_id] = task['time_constraint']

        all_uav_trajectories = {}

        for priority, tasks_in_group in groupby(sorted_tasks, key=lambda x: x['priority']):
            print(f"\n--- Processing Priority Level {priority} ---")

            group_tasks = list(tasks_in_group)
            uavs_in_group = list(set(uav_id for task in group_tasks for uav_id in task['uav_ids']))
            print(f"UAVs in this planning group: {uavs_in_group}")

            for task in group_tasks:
                if task['taskType'] == 'patrol':
                    self._prepare_patrol_goals(task['vertices'], task['uav_ids'], task['time_constraint'])
                elif task['taskType'] == 'waypoint':
                    self._prepare_waypoint_goals(task['goals'], task['uav_ids'])

            approach_trajectories, approach_time_paths = self._execute_planning(
                uavs_in_group,
                task_type=f"priority_{priority}_group", reservation_table=self.reservation_table
            )

            final_trajectories, final_time_paths = self._stitch_patrol_paths(
                approach_trajectories,
                approach_time_paths
            )
            
            for uav_id, path in final_time_paths.items():
                if path:
                    self.reservation_table.reserve_path(path)
            
            all_uav_trajectories.update(final_trajectories)
            self.all_time_paths.update(final_time_paths)

        self._prepare_animation_timestamps()
        return all_uav_trajectories

    def _prepare_animation_timestamps(self):
        """收集并排序所有无人机的时间戳"""
        if not self.all_time_paths: return
        timestamps = set()
        for path in self.all_time_paths.values():
            for _, _, t in path: timestamps.add(round(t, 3))
        self.animation_timestamps = sorted(list(timestamps))
        if self.animation_timestamps:
            print(f"\nAnimation Info: {len(self.animation_timestamps)} frames, Duration: {self.animation_timestamps[-1]:.2f}s")

    def _process_start_positions(self, start_positions: List[Dict]):
        """处理无人机起始位置"""
        for start in start_positions:
            uav_id = start['uav_id']
            if uav_id in self.uav_positions: continue
            pos = start['position']
            lat, lon, alt = pos['lat'], pos['lon'], pos['alt']
            x, y = self.pos_converter.GPStoXY(lat, lon, self.ref_lat, self.ref_lon)
            self.uav_positions[uav_id] = (x, y, alt)
            print(f"UAV {uav_id} start local coordinates: (x={x:.2f}, y={y:.2f})")

    def _prepare_waypoint_goals(self, goals: List[Dict], uav_ids: List[int]):
        """仅准备航点任务的目标点"""
        for goal in goals:
            uav_id = goal['uav_id']
            if uav_id not in uav_ids: continue
            pos = goal['position']
            lat, lon, alt = pos['lat'], pos['lon'], pos['alt']
            x, y = self.pos_converter.GPStoXY(lat, lon, self.ref_lat, self.ref_lon)
            self.uav_goals[uav_id] = (x, y, alt)
            print(f"UAV {uav_id} waypoint goal set to: (x={x:.2f}, y={y:.2f})")

    def _prepare_patrol_goals(self, vertices: List[Dict], uav_ids: List[int], time_constraint: float):
        """准备巡航任务的目标点（入口点）和巡航路径数据"""
        patrol_altitude = self.default_altitude
        if vertices:
            max_alt = max((v['alt'] for v in vertices if v.get('alt', 0.0) != 0.0), default=0.0)
            patrol_altitude = max_alt if max_alt > 0 else self.default_altitude
        
        xy_vertices = [
            {"x": self.pos_converter.GPStoXY(v['lat'], v['lon'], self.ref_lat, self.ref_lon)[0],
             "y": self.pos_converter.GPStoXY(v['lat'], v['lon'], self.ref_lat, self.ref_lon)[1],
             "z": patrol_altitude}
            for v in vertices
        ]
        
        current_patrol_speed = self.patrol_speed
        max_attempts_patrol = 3
        generated_goals = None

        for attempt in range(max_attempts_patrol):
            print(f"Attempt {attempt+1} to generate patrol path with speed: {current_patrol_speed:.2f} m/s")
            generated_goals, patrol_path_points = self._generate_patrol_goals(xy_vertices, uav_ids, patrol_altitude, current_patrol_speed)
            
            if generated_goals:
                break
            
            if current_patrol_speed < self.max_patrol_speed:
                current_patrol_speed = min(current_patrol_speed * self.speed_increase_factor, self.max_patrol_speed)
            else:
                print("Patrol speed reached max, but still could not generate patrol goals.")
                break

        if not generated_goals:
            print(f"Error: Failed to generate patrol goals for UAVs {uav_ids} after all attempts.")
            return

        for goal in generated_goals:
            uav_id = goal["uav_id"]
            entry_point = [goal["position"]["x"], goal["position"]["y"], goal["position"]["z"]]
            path_for_uav = patrol_path_points.get(uav_id, [])
            
            converted_gps_path = [
                {"lat": self.pos_converter.XYtoGPS(p[0], p[1], self.ref_lat, self.ref_lon)[0],
                 "lon": self.pos_converter.XYtoGPS(p[0], p[1], self.ref_lat, self.ref_lon)[1],
                 "alt": patrol_altitude}
                for p in path_for_uav
            ]
            
            total_patrol_length = sum(math.hypot(p2[0]-p1[0], p2[1]-p1[1]) for p1, p2 in zip(path_for_uav[:-1], path_for_uav[1:]))
            estimated_patrol_time = total_patrol_length / current_patrol_speed if current_patrol_speed > 0 else 0

            self.uav_goals[uav_id] = (entry_point[0], entry_point[1], entry_point[2])
            self.cruise_data[uav_id] = {
                "path_points_gps": converted_gps_path,
                "path_points_xy": path_for_uav,
                "estimated_patrol_time": estimated_patrol_time,
                "patrol_speed_used": current_patrol_speed
            }
            self.final_patrol_speeds[uav_id] = current_patrol_speed
            print(f"UAV {uav_id} patrol entry goal set. Estimated patrol time: {estimated_patrol_time:.2f}s")

    def _generate_patrol_goals(self, vertices: List[Dict], uav_ids: List[int], altitude: float, scan_speed: float) -> Tuple[List[Dict], Dict]:
        """调用AreaCoveragePlanner生成巡航任务的目标点和完整巡航路径"""
        try:
            uav_current_positions = [list(self.uav_positions.get(uid, (0, 0, altitude))) for uid in uav_ids]
            polygon_vertices = [[v["x"], v["y"], v["z"]] for v in vertices]
            
            planner = AreaCoveragePlanner(
                uav_ids=uav_ids, polygon_vertices=polygon_vertices, uav_current_positions=uav_current_positions, 
                flight_altitude=altitude, sweep_width=self.sweep_width, scan_speed=scan_speed, 
                turn_speed=scan_speed * 0.5, max_segment_length=self.max_segment_length
            )
            
            path_points, _ = planner.plan_routes()
            
            goals = []
            for uav_id in uav_ids:
                if uav_id in path_points and len(path_points[uav_id]) > 0:
                    first_point = path_points[uav_id][0]
                    goals.append({"uav_id": uav_id, "position": {"x": first_point[0], "y": first_point[1], "z": altitude}})
            
            return goals, path_points
            
        except Exception as e:
            print(f"Area coverage planning failed: {str(e)}. Returning empty goals and paths.")
            return [], {}

    def _stitch_patrol_paths(self, approach_trajectories: Dict, approach_time_paths: Dict) -> Tuple[Dict, Dict]:
        """将接近路径与巡航路径拼接起来"""
        final_trajectories, final_time_paths = {}, {}

        for uav_id, approach_traj in approach_trajectories.items():
            if uav_id in self.cruise_data:
                cruise_info = self.cruise_data[uav_id]
                final_trajectories[uav_id] = approach_traj + cruise_info["path_points_gps"][1:]

                approach_path = approach_time_paths.get(uav_id, [])
                last_time = approach_path[-1][2] if approach_path else 0

                patrol_xy_path = cruise_info["path_points_xy"]
                patrol_speed = cruise_info["patrol_speed_used"]
                
                adjusted_patrol_grid_path = []
                current_time_offset = last_time
                
                for i in range(len(patrol_xy_path)):
                    x, y, _ = patrol_xy_path[i]
                    x_idx = self.fine_planner.calc_xy_index(x, self.fine_planner.min_x)
                    y_idx = self.fine_planner.calc_xy_index(y, self.fine_planner.min_y)
                    adjusted_patrol_grid_path.append((x_idx, y_idx, current_time_offset))
                    
                    if i < len(patrol_xy_path) - 1:
                        dist = math.hypot(patrol_xy_path[i+1][0] - x, patrol_xy_path[i+1][1] - y)
                        current_time_offset += dist / patrol_speed if patrol_speed > 0 else 0
                
                merged_grid_path = list(approach_path)
                if merged_grid_path and adjusted_patrol_grid_path: merged_grid_path.pop()
                merged_grid_path.extend(adjusted_patrol_grid_path)
                final_time_paths[uav_id] = merged_grid_path
            else:
                final_trajectories[uav_id] = approach_traj
                final_time_paths[uav_id] = approach_time_paths.get(uav_id, [])

        return final_trajectories, final_time_paths

    def _execute_planning(self, uav_ids: List[int], task_type: str, reservation_table: ReservationTable) -> Tuple[Dict, Dict]:
        """执行“接近路径”规划，动态调整速度以满足时间约束"""
        current_planning_speed = self.speed
        max_attempts = 5
        
        for attempt in range(max_attempts):
            print(f"\n--- Planning attempt {attempt+1} for {task_type} with speed: {current_planning_speed:.2f} m/s ---")
            
            starts_goals = [(self.uav_positions[uid][:2], self.uav_goals[uid][:2]) for uid in uav_ids if uid in self.uav_positions and uid in self.uav_goals]
            active_uav_ids = [uid for uid in uav_ids if uid in self.uav_positions and uid in self.uav_goals]
            if not starts_goals: return {}, {}

            uav_list_coarse = [(s[0], s[1], g[0], g[1]) for s, g in starts_goals]
            coarse_paths_raw = multi_uav_planning(uav_list_coarse, self.coarse_planner, current_planning_speed, self.max_t, existing_reservation_table=reservation_table)

            if not coarse_paths_raw or any(not p for p in coarse_paths_raw):
                print("Coarse planning failed. Adjusting speed.")
                current_planning_speed = min(current_planning_speed * self.speed_increase_factor, self.max_planning_speed)
                continue
            
            fine_meta = {}
            uav_list_fine = []
            for i, uid in enumerate(active_uav_ids):
                last_x_idx, last_y_idx, last_t = coarse_paths_raw[i][-1]
                fine_start_x = self.coarse_planner.calc_grid_position(last_x_idx, self.coarse_planner.min_x)
                fine_start_y = self.coarse_planner.calc_grid_position(last_y_idx, self.coarse_planner.min_y)
                goal_x, goal_y = self.uav_goals[uid][:2]
                uav_list_fine.append((fine_start_x, fine_start_y, goal_x, goal_y))
                fine_meta[uid] = {'coarse_path': coarse_paths_raw[i], 'end_time': last_t}
            
            fine_paths_raw = multi_uav_planning(uav_list_fine, self.fine_planner, current_planning_speed, self.max_t, existing_reservation_table=reservation_table)        

            if not fine_paths_raw or any(not p for p in fine_paths_raw):
                print("Fine planning failed. Adjusting speed.")
                current_planning_speed = min(current_planning_speed * self.speed_increase_factor, self.max_planning_speed)
                continue
            
            fine_paths_map = {uid: fine_paths_raw[i] for i, uid in enumerate(fine_meta.keys())}
            
            needs_replan, max_time_ratio = self._check_time_constraints(active_uav_ids, fine_meta, fine_paths_map)

            if not needs_replan:
                print(f"\nAll UAVs for {task_type} met time constraints with speed {current_planning_speed:.2f} m/s.")
                return self._finalize_paths(active_uav_ids, fine_meta, fine_paths_map, current_planning_speed)
            else:
                print("Time constraints not met, increasing speed for replanning.")
                new_speed = current_planning_speed * max(self.speed_increase_factor, math.sqrt(max_time_ratio))
                current_planning_speed = min(new_speed, self.max_planning_speed)

        print(f"\nWarning: Could not meet all time constraints after {max_attempts} attempts. Returning best effort result.")
        return self._replan_on_fail(active_uav_ids, current_planning_speed, reservation_table)

    def _check_time_constraints(self, uav_ids, fine_meta, fine_paths_map):
        needs_replan = False
        max_time_ratio = 1.0
        for uid in uav_ids:
            merged_path = self._merge_coarse_fine_paths(fine_meta[uid]['coarse_path'], fine_paths_map[uid], fine_meta[uid]['end_time'])
            path_total_time = merged_path[-1][2] if merged_path else 0
            
            available_time = self.task_time_constraints.get(uid, float('inf')) - self.time_buffer
            if uid in self.cruise_data:
                available_time -= self.cruise_data[uid]["estimated_patrol_time"]
            available_time = max(0.1, available_time)

            print(f"UAV {uid}: Path time: {path_total_time:.2f}s, Available for approach: {available_time:.2f}s")
            if path_total_time > available_time:
                needs_replan = True
                max_time_ratio = max(max_time_ratio, path_total_time / available_time)
        return needs_replan, max_time_ratio

    def _finalize_paths(self, uav_ids, fine_meta, fine_paths_map, speed):
        trajectories, time_paths = {}, {}
        for uid in uav_ids:
            self.final_planning_speeds[uid] = speed
            merged_path = self._merge_coarse_fine_paths(fine_meta[uid]['coarse_path'], fine_paths_map[uid], fine_meta[uid]['end_time'])
            trajectories[uid] = self._convert_path_to_waypoints(merged_path, uid)
            time_paths[uid] = merged_path
        return trajectories, time_paths

    def _merge_coarse_fine_paths(self, coarse_path, fine_path, coarse_end_time):
        """合并粗略和精细路径，统一到精细栅格坐标系"""
        merged_path = []
        for x_idx_c, y_idx_c, t_c in coarse_path:
            x_real = self.coarse_planner.calc_grid_position(x_idx_c, self.coarse_planner.min_x)
            y_real = self.coarse_planner.calc_grid_position(y_idx_c, self.coarse_planner.min_y)
            x_idx_f = self.fine_planner.calc_xy_index(x_real, self.fine_planner.min_x)
            y_idx_f = self.fine_planner.calc_xy_index(y_real, self.fine_planner.min_y)
            merged_path.append((x_idx_f, y_idx_f, t_c))
        
        if fine_path:
            fine_path_offset = [(x, y, coarse_end_time + t) for x, y, t in fine_path]
            if merged_path and merged_path[-1][:2] == fine_path_offset[0][:2]:
                merged_path.pop()
            merged_path.extend(fine_path_offset)
        return merged_path

    def _replan_on_fail(self, uav_ids: List[int], final_speed: float, reservation_table: ReservationTable) -> Tuple[Dict, Dict]:
        """当无法满足约束时，用最终速度再执行一次规划并返回结果"""
        print(f"Executing final planning pass with speed {final_speed:.2f} m/s.")
        starts_goals = [(self.uav_positions[uid][:2], self.uav_goals[uid][:2]) for uid in uav_ids if uid in self.uav_positions and uid in self.uav_goals]
        if not starts_goals: return {}, {}

        # 简单地执行一次规划，不检查时间约束
        uav_list_coarse = [(s[0], s[1], g[0], g[1]) for s, g in starts_goals]
        coarse_paths_raw = multi_uav_planning(uav_list_coarse, self.coarse_planner, final_speed, self.max_t, existing_reservation_table=reservation_table)
        if not coarse_paths_raw or any(not p for p in coarse_paths_raw): return {}, {}
            
        fine_meta = {}
        uav_list_fine = []
        for i, uid in enumerate(uav_ids):
            last_x_idx, last_y_idx, last_t = coarse_paths_raw[i][-1]
            fine_start_x = self.coarse_planner.calc_grid_position(last_x_idx, self.coarse_planner.min_x)
            fine_start_y = self.coarse_planner.calc_grid_position(last_y_idx, self.coarse_planner.min_y)
            goal_x, goal_y = self.uav_goals[uid][:2]
            uav_list_fine.append((fine_start_x, fine_start_y, goal_x, goal_y))
            fine_meta[uid] = {'coarse_path': coarse_paths_raw[i], 'end_time': last_t}
            
        fine_paths_raw = multi_uav_planning(uav_list_fine, self.fine_planner, final_speed, self.max_t, existing_reservation_table=reservation_table)
        if not fine_paths_raw or any(not p for p in fine_paths_raw): return {}, {}

        fine_paths_map = {uid: fine_paths_raw[i] for i, uid in enumerate(fine_meta.keys())}
        return self._finalize_paths(uav_ids, fine_meta, fine_paths_map, final_speed)

    def _convert_path_to_waypoints(self, path: List[Tuple[int, int, float]], uav_id: int) -> List[Dict]:
        """将栅格路径转换为经纬度航点列表"""
        waypoints = []
        alt = self.uav_goals.get(uav_id, (0, 0, self.default_altitude))[2]
        for x_idx, y_idx, _ in path:
            x = self.fine_planner.calc_grid_position(x_idx, self.fine_planner.min_x)
            y = self.fine_planner.calc_grid_position(y_idx, self.fine_planner.min_y)
            lat, lon = self.pos_converter.XYtoGPS(x, y, self.ref_lat, self.ref_lon)
            waypoints.append({"lat": lat, "lon": lon, "alt": alt})
        return waypoints


# --- 服务层入口函数 ---
def run_planning_service(request_data: Dict) -> Tuple[Dict, Dict]:
    """
    服务层的主函数，被API层调用。
    :param request_data: 包含 'tasks' 和 'config' 的字典。
    :return: 一个包含轨迹字典和时空路径字典的元组。
    """
    start_time = datetime.now()
    logger.info("="*60)
    logger.info("[规划服务] 开始执行规划任务")
    logger.info(f"[规划服务] 任务数量: {len(request_data.get('tasks', []))}")
    logger.info(f"[规划服务] 配置信息: {request_data.get('config', {})}")
    
    tasks = request_data['tasks']
    config_dict = request_data['config']
    
    # 使用 Pydantic 模型验证和创建配置对象
    config = PlannerConfig.model_validate(config_dict)
    logger.info(f"[规划服务] 配置验证完成")
    
    # 1. 创建规划器实例，并传入配置
    logger.info("[规划服务] 创建规划器实例...")
    planner = MultiUAVPlanner(config)
    
    # 2. 执行规划
    logger.info("[规划服务] 开始执行任务规划...")
    planning_start = datetime.now()
    trajectories_result = planner.process_tasks(tasks)
    planning_end = datetime.now()
    logger.info(f"[规划服务] 任务规划完成，耗时: {(planning_end - planning_start).total_seconds():.2f}秒")
    
    # 3. 提取最终的时空路径数据以供返回
    time_paths_result = planner.all_time_paths
    
    logger.info(f"[规划服务] 提取时空路径数据，路径数量: {len(time_paths_result)}")
    
    # 转换为API模型要求的格式: Dict[str, List[List[Union[int, float]]]]
    #time_paths已经是 (int, int, float)的元组列表，需要转换成列表的列表
    formatted_time_paths = {str(k): [list(point) for point in v] for k, v in time_paths_result.items()}

    #转换轨迹字典的key为字符串，以符合Pydantic模型
    formatted_trajectories = {
        str(uav_id): {
            "trajectory": traj,
            "suggested_planning_speed": planner.final_planning_speeds.get(uav_id),
            "suggested_patrol_speed": planner.final_patrol_speeds.get(uav_id),
        }
        for uav_id, traj in trajectories_result.items()
    }
    
    end_time = datetime.now()
    total_time = (end_time - start_time).total_seconds()
    
    logger.info(f"[规划服务] 规划任务完成")
    logger.info(f"[规划服务] 生成轨迹数量: {len(formatted_trajectories)}")
    logger.info(f"[规划服务] 生成路径数量: {len(formatted_time_paths)}")
    logger.info(f"[规划服务] 总耗时: {total_time:.2f}秒")
    logger.info("="*60)
    
    return formatted_trajectories, formatted_time_paths
